import pandas as pd
from tqdm import tqdm
def gen_voc():
    with open('E:/CVFO/yan/simplebooks_92/train.txt', 'r', encoding="utf-8") as f:
        lines = f.readlines()
    train_data = [i.strip().split() for i in tqdm(lines) if len(i) > 2 and len(i.strip().split()) < 136]

    with open('E:/CVFO/yan/simplebooks_92/valid.txt', 'r', encoding="utf-8") as f:
        lines = f.readlines()
    val_data = [i.strip().split() for i in tqdm(lines) if len(i) > 2 and len(i.strip().split()) < 136]

    with open('E:/CVFO/yan/simplebooks_92/test.txt', 'r', encoding="utf-8") as f:
        lines = f.readlines()
    test_data = [i.strip().split() for i in tqdm(lines) if len(i) > 2 and len(i.strip().split()) < 136]

    pos=["<|p_{}|>".format(i) for i in range(139)]
    voc=dict()
    for data in tqdm(train_data+val_data+test_data):
        for word in data:
            if voc.get(word,None) is None:
                voc[word]=1
    pd.to_pickle(["<|pad|>","<|aos|>","<|bos|>"]+list(voc.keys())+pos,"E:/CVFO/yan/simplebooks_92/vocab.pkl")


gen_voc()

